// Copyright (c) 2024 Go-Frame-Lite
// Go-Frame-Lite is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan
// PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//         http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

package middleware

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"io"
	"net/http"

	"go-frame-lite/pkg/crypto"
	"go-frame-lite/pkg/logger"

	"github.com/gin-gonic/gin"
	"github.com/spf13/viper"
	"go.uber.org/zap"
)

// SM4EncryptionConfig SM4数据加密配置
type SM4EncryptionConfig struct {
	Enabled          bool   `mapstructure:"enabled"`
	Key              string `mapstructure:"key"`
	Mode             string `mapstructure:"mode"`
	EncryptionHeader string `mapstructure:"encryption_header"`
}

// responseWriter 自定义响应写入器，用于捕获响应体
type responseWriter struct {
	gin.ResponseWriter
	body *bytes.Buffer
}

// Write 重写Write方法，将响应体写入缓冲区
func (w responseWriter) Write(b []byte) (int, error) {
	w.body.Write(b)
	return w.ResponseWriter.Write(b)
}

// SM4EncryptionMiddleware SM4数据加密中间件
func SM4EncryptionMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 获取配置
		config := getSM4Config()
		
		// 如果未启用SM4加密，直接跳过
		if !config.Enabled {
			c.Next()
			return
		}

		// 检查是否为加密请求
		if isEncryptedRequest(c, config) {
			// 解密请求体
			if err := decryptRequestBody(c, config); err != nil {
				logger.Error(c.Request.Context(), "SM4解密失败", 
					zap.String("error", err.Error()))
				c.JSON(http.StatusBadRequest, gin.H{
					"error": "解密失败",
				})
				c.Abort()
				return
			}
		}

		// 如果需要加密响应，创建自定义响应写入器
		var writer gin.ResponseWriter
		var bodyBuffer *bytes.Buffer
		if shouldEncryptResponse(c, config) {
			bodyBuffer = &bytes.Buffer{}
			writer = &responseWriter{
				ResponseWriter: c.Writer,
				body:          bodyBuffer,
			}
			c.Writer = writer
		}

		// 处理请求
		c.Next()

		// 加密响应体
		if shouldEncryptResponse(c, config) && bodyBuffer != nil && bodyBuffer.Len() > 0 {
			if err := encryptResponseBody(c, config, bodyBuffer); err != nil {
				logger.Error(c.Request.Context(), "SM4加密失败", 
					zap.String("error", err.Error()))
				// 如果加密失败，保留原始响应
			}
		}
	}
}

// getSM4Config 获取SM4配置
func getSM4Config() SM4EncryptionConfig {
	return SM4EncryptionConfig{
		Enabled:          viper.GetBool("crypto.sm4.enabled"),
		Key:              viper.GetString("crypto.sm4.key"),
		Mode:             viper.GetString("crypto.sm4.mode"),
		EncryptionHeader: "X-SM4-Encrypted",
	}
}

// isEncryptedRequest 检查是否为加密请求
func isEncryptedRequest(c *gin.Context, config SM4EncryptionConfig) bool {
	encryptionHeader := c.GetHeader(config.EncryptionHeader)
	return encryptionHeader == "true" || encryptionHeader == "1"
}

// shouldEncryptResponse 检查是否需要加密响应
func shouldEncryptResponse(c *gin.Context, config SM4EncryptionConfig) bool {
	// 根据请求头或配置决定是否加密响应
	encryptionHeader := c.GetHeader(config.EncryptionHeader)
	return encryptionHeader == "true" || encryptionHeader == "1"
}

// decryptRequestBody 解密请求体
func decryptRequestBody(c *gin.Context, config SM4EncryptionConfig) error {
	// 读取原始请求体
	body, err := io.ReadAll(c.Request.Body)
	if err != nil {
		return fmt.Errorf("读取请求体失败: %w", err)
	}

	// 如果请求体为空，无需解密
	if len(body) == 0 {
		return nil
	}

	// Base64解码
	decodedBody, err := base64.StdEncoding.DecodeString(string(body))
	if err != nil {
		return fmt.Errorf("Base64解码失败: %w", err)
	}

	// 解码并验证密钥
	key, err := decodeAndValidateSM4Key(config.Key)
	if err != nil {
		return err
	}

	// 使用项目中的crypto包进行解密
	decryptedData, err := crypto.SM4Decrypt(key, decodedBody)
	if err != nil {
		return fmt.Errorf("SM4解密失败: %w", err)
	}

	// 替换请求体
	c.Request.Body = io.NopCloser(bytes.NewReader(decryptedData))
	c.Request.ContentLength = int64(len(decryptedData))

	return nil
}

// encryptResponseBody 加密响应体
func encryptResponseBody(c *gin.Context, config SM4EncryptionConfig, bodyBuffer *bytes.Buffer) error {
	// 获取响应体
	bodyBytes := bodyBuffer.Bytes()

	// 如果响应体为空，无需加密
	if len(bodyBytes) == 0 {
		return nil
	}

	// 解码并验证密钥
	key, err := decodeAndValidateSM4Key(config.Key)
	if err != nil {
		return err
	}

	// 使用项目中的crypto包进行加密
	encryptedData, err := crypto.SM4Encrypt(key, bodyBytes)
	if err != nil {
		return fmt.Errorf("SM4加密失败: %w", err)
	}

	// Base64编码
	encodedData := base64.StdEncoding.EncodeToString(encryptedData)

	// 清除之前的响应头和内容
	c.Header("Content-Type", "text/plain")
	c.Header("Content-Length", fmt.Sprintf("%d", len(encodedData)))
	c.Writer.WriteHeaderNow()

	// 写入加密后的响应体
	_, err = c.Writer.Write([]byte(encodedData))
	return err
}

// decodeAndValidateSM4Key 解码并验证SM4密钥
func decodeAndValidateSM4Key(keyStr string) ([]byte, error) {
	// Base64解码密钥
	key, err := base64.StdEncoding.DecodeString(keyStr)
	if err != nil {
		return nil, fmt.Errorf("密钥Base64解码失败: %w", err)
	}

	// 验证密钥长度是否为16字节
	if len(key) != 16 {
		return nil, fmt.Errorf("SM4密钥长度必须为16字节，当前长度: %d", len(key))
	}

	return key, nil
}